#pragma once
// -------------------------------------------------
// ZED Kit
// -------------------------------------------------
//   File Name: array.hpp
//      Author: Ziming Li
//     Created: 2023-06-04
// -------------------------------------------------
// Copyright (C) 2023 MingYang Software Technology.
// -------------------------------------------------

#ifndef ZED_ARRAY_HPP
#define ZED_ARRAY_HPP

#include "./assert.h"
#include <utility>
#ifndef NDEBUG
#   include <vector>
#endif

namespace zed {

template <typename T, bool auto_construct = true>
class flexible_array
{
public:
    static flexible_array* create(size_t size);
    ~flexible_array(void);

    void* operator new(size_t, void* ptr) { return ptr; }
    void operator delete(void* p) { free(p); }

    template <typename... Args>
    std::enable_if_t<!auto_construct> construct(size_t i, Args&&... args);

    T& operator[](size_t i);
    const T& operator[](size_t i) const;

    size_t size(void) const { return m_size; }

    template <bool>
    class iterator_impl;

    using iterator = iterator_impl<false>;
    iterator begin(void);
    iterator end(void);

    using const_iterator = iterator_impl<true>;
    const_iterator begin(void) const;
    const_iterator end(void) const;
protected:
    flexible_array(size_t size);

    static size_t calculate_size(size_t n) { return sizeof(flexible_array) + sizeof(T) * n; }
private:
    void* operator new(size_t size) { return malloc(size); }

    T* get_raw_array(void) const { return reinterpret_cast<T *>(const_cast<flexible_array*>(this) + 1); }

    const size_t m_size;
#ifndef NDEBUG
    std::vector<bool> m_construct_flags;
#endif
};

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Implementations

template <typename T, bool auto_construct>
template <bool is_const>
class flexible_array<T, auto_construct>::iterator_impl
{
    friend class flexible_array<T, auto_construct>;
public:
    using value_type = std::conditional_t<is_const, const T, T>;

    bool operator==(const iterator_impl& o) const { return m_current == o.m_current; }

    iterator_impl& operator++()
    {
        ++m_current;
        return *this;
    }

    value_type& operator*() const { return *m_current; }
private:
    iterator_impl(value_type* current) : m_current(current) {}

    value_type* m_current;
};

template <typename T, bool auto_construct>
flexible_array<T, auto_construct>::flexible_array(size_t size) : m_size(size)
{
    if (0 == m_size)
        return;

    if constexpr (auto_construct)
    {
        T* raw_array = get_raw_array();
        new (raw_array) T[m_size];
    }
#ifndef NDEBUG
    m_construct_flags.resize(m_size, auto_construct);
#endif
}

template <typename T, bool auto_construct>
flexible_array<T, auto_construct>::~flexible_array(void)
{
    if (0 == m_size)
        return;

    T* raw_array = get_raw_array();
    for (size_t i = 0; i < m_size; ++i)
    {
        ZASSERT(m_construct_flags[i]);
        raw_array[i].~T();
    }
}

template <typename T, bool auto_construct>
flexible_array<T, auto_construct>* flexible_array<T, auto_construct>::create(size_t size) {
    void* p = malloc(calculate_size(size));
    return new (p) flexible_array(size);
}

template <typename T, bool auto_construct>
template <typename... Args>
std::enable_if_t<!auto_construct> flexible_array<T, auto_construct>::construct(size_t i, Args&&... args)
{
    ZASSERT(i < m_size && !m_construct_flags[i]);
    new (get_raw_array() + i) T(std::forward<Args>(args)...);
#ifndef NDEBUG
    m_construct_flags[i] = true;
#endif
}

template <typename T, bool auto_construct>
T& flexible_array<T, auto_construct>::operator[](size_t i) {
    ZASSERT(i < m_size && m_construct_flags[i]);
    return get_raw_array()[i];
}

template <typename T, bool auto_construct>
const T& flexible_array<T, auto_construct>::operator[](size_t i) const {
    ZASSERT(i < m_size && m_construct_flags[i]);
    return get_raw_array()[i];
}

template <typename T, bool auto_construct>
typename flexible_array<T, auto_construct>::iterator flexible_array<T, auto_construct>::begin(void)
{
    return iterator(get_raw_array());
}

template <typename T, bool auto_construct>
typename flexible_array<T, auto_construct>::iterator flexible_array<T, auto_construct>::end(void)
{
    return iterator(get_raw_array() + m_size);
}

template <typename T, bool auto_construct>
typename flexible_array<T, auto_construct>::const_iterator flexible_array<T, auto_construct>::begin(void) const
{
    return const_iterator(get_raw_array());
}

template <typename T, bool auto_construct>
typename flexible_array<T, auto_construct>::const_iterator flexible_array<T, auto_construct>::end(void) const
{
    return const_iterator(get_raw_array() + m_size);
}

} // namespace zed

#endif // ZED_ARRAY_HPP
